package com.gitee.l0km.xthrift.service.mq;

import java.util.concurrent.Executor;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import org.apache.thrift.TApplicationException;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TMessage;
import org.apache.thrift.protocol.TMessageType;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.AutoExpandingBufferWriteTransport;
import org.apache.thrift.transport.TMemoryInputTransport;
import org.jboss.netty.util.Timeout;
import org.jboss.netty.util.Timer;
import org.jboss.netty.util.TimerTask;

import com.facebook.nifty.core.RequestContext;
import com.facebook.nifty.core.RequestContexts;
import com.facebook.nifty.duplex.TDuplexProtocolFactory;
import com.facebook.nifty.duplex.TProtocolPair;
import com.facebook.nifty.duplex.TTransportPair;
import com.facebook.nifty.processor.NiftyProcessor;
import com.gitee.l0km.com4j.base.SimpleLog;
import com.gitee.l0km.xthrift.mq.Envelope;
import com.gitee.l0km.xthrift.service.XThriftRequestContext;
import com.gitee.l0km.xthrift.service.XthriftConstant;
import com.gitee.l0km.xthrift.thrift.protocol.TJSONProtocolXthrift;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;

import gu.simplemq.Channel;
import gu.simplemq.IMessageAdapter;
import gu.simplemq.IPublisher;
import gu.simplemq.exceptions.SmqUnsubscribeException;
import io.airlift.units.Duration;

import static com.google.common.base.Strings.isNullOrEmpty;
import static com.google.common.base.MoreObjects.firstNonNull;

/**
 * 基于消息系统（simplemq）的RPC请求分发实现 
 * @author guyadong
 * @since 1.4.0
 *
 */
public class SimplemqDispatcher implements IMessageAdapter<Envelope>,XthriftConstant {
	private NiftyProcessor processor;
	private TDuplexProtocolFactory duplexProtocolFactory;
	private Executor executor;
	private long taskTimeoutMillis;
	private Timer taskTimeoutTimer;
	private IPublisher publisher;
	public SimplemqDispatcher(NiftyProcessor processor,
			Executor executor,
			Duration taskTimeout,
			Timer timer,
			IPublisher publisher) {
		this.processor = processor;
		this.duplexProtocolFactory = TDuplexProtocolFactory.fromSingleFactory(new TJSONProtocolXthrift.Factory(true, true, false));
		this.executor = executor;
		this.taskTimeoutMillis = (taskTimeout == null ? 0 : taskTimeout.toMillis());
		this.taskTimeoutTimer = (taskTimeout == null ? null : timer);
		this.publisher = publisher;
	}

	public Channel<Envelope> channelOf(String name){
		return  new Channel<>(name, Envelope.class).setAdapter(this);
	}

	@Override
	public void onSubscribe(Envelope envelope) throws SmqUnsubscribeException {
		/** 忽略请求发起时间超时的请求 */
		if (null != envelope.getSendTime() && Math.abs(System.currentTimeMillis() - envelope.getSendTime().getTime()) > taskTimeoutMillis) {
			return;
		}
		if (taskTimeoutMillis > 0) {
			envelope.processStartTimeMillis(System.currentTimeMillis());
		}
		TMemoryInputTransport inputTransport = new TMemoryInputTransport(envelope.getBody().getBytes());
		AutoExpandingBufferWriteTransport outputTransport = new AutoExpandingBufferWriteTransport(1024, 1.5);
		TTransportPair transportPair = TTransportPair.fromSeparateTransports(inputTransport, outputTransport);
		TProtocolPair protocolPair = duplexProtocolFactory.getProtocolPair(transportPair);
		TProtocol inProtocol = protocolPair.getInputProtocol();
		TProtocol outProtocol = protocolPair.getOutputProtocol();
		RequestContext requestContext = new XThriftRequestContext(null, inProtocol, outProtocol);
		if(null != envelope.getThinError()) {
			requestContext.setContextData(CTX_THIN_ERROR, envelope.getThinError());
		}
		processRequest(envelope, requestContext, inProtocol, outProtocol);
	}
    private void processRequest(
            final Envelope envelope,
            final RequestContext requestContext,
            final TProtocol inProtocol, 
            final TProtocol outProtocol) {

        try {
            executor.execute(new Runnable() {
                @Override
                public void run() {
                    ListenableFuture<Boolean> processFuture;
                    final AtomicBoolean responseSent = new AtomicBoolean(false);
                    // Use AtomicReference as a generic holder class to be able to mark it final
                    // and pass into inner classes. Since we only use .get() and .set(), we don't
                    // actually do any atomic operations.
                    final AtomicReference<Timeout> expireTimeout = new AtomicReference<>(null);

                    try {
                        try {
                            long timeRemaining = 0;
                            if (taskTimeoutMillis > 0) {
                                long timeElapsed = System.currentTimeMillis() - envelope.processStartTimeMillis();
                                if (timeElapsed >= taskTimeoutMillis) {
                                    TApplicationException taskTimeoutException = new TApplicationException(
                                            TApplicationException.INTERNAL_ERROR,
                                            "Task stayed on the queue for " + timeElapsed +
                                                    " milliseconds, exceeding configured task timeout of " + taskTimeoutMillis +
                                                    " milliseconds."
                                    );
                                    sendTApplicationException(taskTimeoutException, envelope, inProtocol, outProtocol);
                                    return;
                                } else {
                                    timeRemaining = taskTimeoutMillis - timeElapsed;
                                }
                            }

                            if (timeRemaining > 0) {
                                expireTimeout.set(taskTimeoutTimer.newTimeout(new TimerTask() {
                                    @Override
                                    public void run(Timeout timeout) throws Exception {
                                        // The immediateFuture returned by processors isn't cancellable, cancel() and
                                        // isCanceled() always return false. Use a flag to detect task expiration.
                                        if (responseSent.compareAndSet(false, true)) {
                                            TApplicationException ex = new TApplicationException(
                                                    TApplicationException.INTERNAL_ERROR,
                                                    "Task timed out while executing."
                                            );
											sendTApplicationException(ex, envelope, inProtocol, outProtocol);
                                        }
                                    }
                                }, timeRemaining, TimeUnit.MILLISECONDS));
                            }

                            RequestContexts.setCurrentContext(requestContext);
                            processFuture = processor.process(inProtocol, outProtocol, requestContext);
                        } finally {
                            // RequestContext does NOT stay set while we are waiting for the process
                            // future to complete. This is by design because we'll might move on to the
                            // next request using this thread before this one is completed. If you need
                            // the context throughout an asynchronous handler, you need to read and store
                            // it before returning a future.
                            RequestContexts.clearCurrentContext();
                        }

                        Futures.addCallback(
                                processFuture,
                                new FutureCallback<Boolean>() {
                                    @Override
                                    public void onSuccess(@Nullable Boolean result) {
                                        deleteExpirationTimer(expireTimeout.get());
                                        try {
                                            // Only write response if the client is still there and the task timeout
                                            // hasn't expired.
                                            if (responseSent.compareAndSet(false, true)) {
                                                writeResponse(envelope, outProtocol);
                                            }
                                        } catch (Throwable t) {
                                            onDispatchException(t);
                                        }
                                    }

                                    @Override
                                    public void onFailure(@Nonnull Throwable t) {
                                        deleteExpirationTimer(expireTimeout.get());
                                        onDispatchException(t);
                                    }
                                }
                        );
                    } catch (TException e) {
                        onDispatchException(e);
                    }
                }
            });
        }
        catch (RejectedExecutionException ex) {
            TApplicationException x = new TApplicationException(TApplicationException.INTERNAL_ERROR,
                    "Server overloaded");
            sendTApplicationException(x, envelope, inProtocol, outProtocol);
        }
    }
    private void deleteExpirationTimer(Timeout timeout)
    {
        if (timeout == null) {
            return;
        }
        timeout.cancel();
    }
    private void sendTApplicationException(
            TApplicationException x,
            Envelope envelope,
            TProtocol inProtocol,
            TProtocol outProtocol)
    {
        try {
        	TMessage message = inProtocol.readMessageBegin();
        	outProtocol.writeMessageBegin(new TMessage(message.name, TMessageType.EXCEPTION, message.seqid));
        	x.write(outProtocol);
        	outProtocol.writeMessageEnd();
        	outProtocol.getTransport().flush();
        	
        	writeResponse(envelope, outProtocol);
        }
        catch (TException ex) {
        	onDispatchException(ex);
        }
    }
    private void onDispatchException(Throwable t)
    {
    	SimpleLog.log(firstNonNull(t.getMessage(),t.getClass().getSimpleName()),t);
    }
	private void writeResponse(Envelope envelope, TProtocol outProtocol) {
		if(!isNullOrEmpty(envelope.getAck())) {
			Channel<Envelope> ch = new Channel<>(envelope.getAck(), Envelope.class);
			AutoExpandingBufferWriteTransport outputTransport = (AutoExpandingBufferWriteTransport)outProtocol.getTransport();
			Envelope env = Envelope.create(outputTransport.getBuf().array(), 0, outputTransport.getPos());
			publisher.publish(ch,env);
		}
	}
}
